from matplotlib.pyplot import figure, show
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import numpy as np

# load files as averages
dt0 = np.average(np.loadtxt('overlap[dt=0].out'), 0)

dtplus10   = np.average(np.loadtxt('overlap[dt=10].out'), 0)
dtminus10  = np.average(np.loadtxt('overlap[dt=-10].out'), 0)
dtplus25   = np.average(np.loadtxt('overlap[dt=25].out'), 0)
dtminus25  = np.average(np.loadtxt('overlap[dt=-25].out'), 0)
dtminus50  = np.average(np.loadtxt('overlap[dt=-50].out'), 0)
dtplus50   = np.average(np.loadtxt('overlap[dt=50].out'), 0)
dtplus100  = np.average(np.loadtxt('overlap[dt=100].out'), 0)
dtminus100 = np.average(np.loadtxt('overlap[dt=-100].out'), 0)


fig = figure()
ax = fig.add_subplot(111, projection='3d')
dt = np.array([-50, -25, 10, 0, 10, 25])

X, Y = np.meshgrid(dt, np.arange(11))
Z = np.ones((11, len(dt)))*0.3
#Z[0] = dtminus100
#Z[1] = dtminus50

z = np.array([dtminus50, dtminus25, dtminus10, dt0, dtplus10, dtplus25])

print(X.shape)
print(Y.shape)
print(Z.shape)

#Z = np.meshgrid(z,z)
#surf = ax.plot_surface(X, Y,  z.T, cstride=1, rstride=1, lw=0.5, cmap=cm.jet, alpha=0.3, antialiased=True)
ax.plot_surface(X, Y,  z.T, cstride=1, rstride=1, lw=0.5, cmap=cm.jet, alpha=0.3, antialiased=True)
ax.plot_wireframe(X, Y,  z.T, cstride=1, rstride=1, color='black')
ax.set_zlim(0, 1.2)
ax.set_xlim(-25, 25)
ax.set_ylim(-2, 11)
ax.set_ylabel('Y Recall step')
ax.set_xlabel('X $\Delta t$')
ax.set_zlabel('Overlap')

#fig.colorbar(surf, shrink=0.5, aspect=5, ticks=np.linspace(0,1,5))
show()
